import org.bytedeco.javacv.*;
import org.bytedeco.opencv.opencv_core.Mat;
import org.bytedeco.opencv.opencv_videoio.VideoWriter;

import java.io.File;
import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.concurrent.TimeUnit;

import static org.bytedeco.opencv.global.opencv_imgcodecs.imwrite;
import static org.bytedeco.opencv.global.opencv_videoio.*;

/**
 * 使用JavaCV处理RTSP流，抽帧并以10fps保存视频
 */
public class RtspFrameProcessor {
    // 目标帧率：10fps
    private static final double TARGET_FPS = 10;
    // 每帧之间的间隔时间(毫秒)
    private static final long FRAME_INTERVAL_MS = (long)(1000 / TARGET_FPS);
    
    private final String rtspUrl;
    private final String outputDir;
    private volatile boolean isRunning = true;

    public RtspFrameProcessor(String rtspUrl, String outputDir) {
        this.rtspUrl = rtspUrl;
        this.outputDir = outputDir;
        // 创建输出目录
        new File(outputDir).mkdirs();
    }

    /**
     * 开始处理RTSP流
     */
    public void startProcessing() {
        FrameGrabber grabber = null;
        FrameRecorder recorder = null;
        
        try {
            // 创建RTSP抓取器
            grabber = FrameGrabber.createDefault(rtspUrl);
            // 设置RTSP传输方式为TCP
            grabber.setOption("rtsp_transport", "tcp");
            // 设置超时时间
            grabber.setOption("stimeout", "5000000"); // 5秒
            grabber.start();
            
            // 获取视频基本信息
            int width = grabber.getImageWidth();
            int height = grabber.getImageHeight();
            double sourceFps = grabber.getFrameRate();
            
            System.out.println("视频信息 - 宽度: " + width + ", 高度: " + height + ", 原帧率: " + sourceFps);
            
            // 创建输出视频文件
            String timestamp = new SimpleDateFormat("yyyyMMdd_HHmmss").format(new Date());
            String outputFile = outputDir + File.separator + "output_" + timestamp + ".mp4";
            
            // 创建视频录制器
            recorder = FrameRecorder.createDefault(outputFile, width, height);
            recorder.setFrameRate(TARGET_FPS);
            recorder.setVideoCodec(avcodec.AV_CODEC_ID_H264);
            recorder.setFormat("mp4");
            recorder.start();
            
            Frame frame;
            long lastFrameTime = 0;
            
            System.out.println("开始处理RTSP流，输出文件: " + outputFile);
            
            // 循环抓取帧
            while (isRunning && (frame = grabber.grab()) != null) {
                long currentTime = System.currentTimeMillis();
                
                // 控制帧率：只保留间隔足够的帧
                if (currentTime - lastFrameTime >= FRAME_INTERVAL_MS) {
                    // 记录当前帧时间
                    lastFrameTime = currentTime;
                    
                    // 写入视频
                    recorder.record(frame);
                    
                    // 可选：输出处理信息
                    // System.out.println("处理帧 - 时间: " + currentTime);
                }
                
                // 短暂休眠，降低CPU占用
                TimeUnit.MILLISECONDS.sleep(1);
            }
            
            System.out.println("处理完成");
            
        } catch (Exception e) {
            System.err.println("处理过程中发生错误: " + e.getMessage());
            e.printStackTrace();
        } finally {
            // 释放资源
            try {
                if (recorder != null) {
                    recorder.stop();
                    recorder.release();
                }
                if (grabber != null) {
                    grabber.stop();
                    grabber.release();
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 停止处理
     */
    public void stopProcessing() {
        isRunning = false;
    }

    public static void main(String[] args) {
        // 示例配置
        String rtspUrl = "rtsp://example.com/stream"; // 替换为实际RTSP地址
        String outputDir = "./output_videos";
        
        RtspFrameProcessor processor = new RtspFrameProcessor(rtspUrl, outputDir);
        
        // 启动处理线程
        Thread processingThread = new Thread(processor::startProcessing);
        processingThread.start();
        
        // 示例：运行30秒后停止
        try {
            TimeUnit.SECONDS.sleep(30);
            processor.stopProcessing();
            processingThread.join();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }
}
